#Author: Mike Schladt - 2014
#Malware Intelligence Project MalwareReport Class
#Defines structures require to process cuckoo reports

import os
import time
import requests
import json
import sets
from termcolor import colored
import logging
logging.basicConfig(filename='mip.log',level=logging.WARNING)

try :
    import cPickle as pickle
except :
    import pickle 

def process_cuckoo_report(sample):

    sid = sample.sid
    task_id = sample.task_id
    #wait until status is reported
    while 1 :
        status = get_task_status(task_id)
        print "Sample: " + str(sid) + " Task: " + str(task_id) + ": " + status 
        if status == 'reported':
            try:
                json_report = get_report(task_id)
                tags = create_tags(json_report)    
                store_tags(sid,tags)
                sample.errors = check_errors(json_report)
            except Exception as e:
                sample.errors = True
                print colored("Error processing report {0} for sample {1}: {2}".format(task_id , sid , e),"red")
                logging.warning("Error processing report {0} for sample {1}: {2}".format(task_id , sid , e))
            if sample.errors == True :
                print colored("Cuckoo Reported Errors for Task {0}".format(task_id),"red")
                print colored("Saving Sample {0} for Re-RUN".format(sid),"red")
                sample_error(sample)
            else :
                print colored("No Errors for Task {0}".format(task_id),"green")
                print colored("Removing Sample {0}".format(sid),"green")
                try :
                    os.remove(sample.file)
                except Exception as e :
                    msg = "Error removing file: {0}".format(e)
                    print msg
                    logging.warning(msg)                      
            return

        else :
            time.sleep(6)

#function to handle samples that error out
#move to folder for further analysis 
def sample_error(sample):
    error_dir = os.environ['HOME']+"/samples/errors/"
    if not os.path.exists(error_dir):
        os.makedirs(error_dir)
    try:
        src = sample.file
        dst = os.path.join(error_dir, os.path.basename(sample.file))
        os.rename(src, dst)
    except Exception as e:
        msg = "Error saving to errors folder: {0}".format(e)
        print msg
        logging.warning(msg)

#module to return status
def get_task_status(task_id) :
    request_url = "http://localhost:8090/tasks/view/" + str(task_id)      
    response = requests.get(request_url)
    status = response.json()['task']['status']
    return status

#module to retrive report
def get_report(task_id) :
    request_url = "http://localhost:8090/tasks/report/" + str(task_id)      
    response = requests.get(request_url)
    report = response.json()
    return report


#returns the cuckoo system status
def get_cuckoo_status() : 
    request_url = "http://localhost:8090/cuckoo/status"      
    response = requests.get(request_url).json()
    return response

#module to create tags from json file
def create_tags(json_report) :
    #tags set to eliminate duplicates
    tags = set()

    #let's start with strings - they're easy
    try:    
        for report_string in json_report['strings'] :
            tag = {}
            tag['type'] = 'string'
            tag['value'] = report_string
            tag['description'] = '[STRING]'
            tags.add(pickle.dumps(tag))
    except e:
        print e
    
    #compile time
    try : 
        tag = {}
        tag['type'] = 'compile_time'
        tag['value'] = json_report['compiletime']
        tag['description'] = 'PE COMPILE TIME: YYYY-MM-DD HH:MM:SS'
        tags.add(pickle.dumps(tag))
    except: pass

    #imported dll count
    try :
        tag = {}
        tag['type'] = 'imported _dll_count'
        tag['value'] = json_report['static']['imported_dll_count']
        tag['description'] = '[NUMBER OF DLLS IMPORTED]'
        tags.add(pickle.dumps(tag))
    except: pass
    
    #import tags
    try :
        for dll in json_report['static']['pe_imports'] :
            tag = {}
            tag['type'] = 'import'
            tag['description'] = '[DLL NAME] : [IMPORT NAME]'
            for pe_import in dll['imports'] :
                tag['value'] = dll['dll'] + ' : ' + pe_import['name']
                tags.add(pickle.dumps(tag))
    except:
        pass    

    #peid_signatures
    try:
        for peid_signature in json_report['static']['peid_signatures'] :
            tag = {}
            tag['type'] = 'peid_signature'
            tag['description'] = '[PEID SIGNARUTE]'
            tag['value'] = peid_signature
            tags.add(pickle.dumps(tag))
    except:
        pass
     
    try:
        for pe_export in json_report['static']['pe_exports'] :
            tag = {}
            tag['type'] = 'export'
            tag['description'] = '[EXPORT NAME]'
            
            tag['value'] = pe_export['name']           
            tags.add(pickle.dumps(tag))
    except:
        pass
   
    # store the PE info
    try: 
        for pe_info in json_report['static']['pe_versioninfo'] :
            tag = {}
            tag['type'] = 'pe_info'
            tag['value'] = pe_info['name'] + ' : ' + pe_info['value']
            tag['description'] = '[NAME] : [VALUE]'
            tags.add(pickle.dumps(tag))
    except:
        pass        

    #section name
    try:
        for section in json_report['static']['pe_sections'] :
            tag = {}
            tag['type'] = 'section_name'
            tag['description'] = '[SECTION NAME]'
            tag['value'] = section['name']
            tags.add(pickle.dumps(tag))
    except: 
        pass

    #section
    try: 
        for section in json_report['static']['pe_sections']:
            tag = {}
            tag['type'] = 'section'
            tag['description'] = "[NAME] : [VIRTUAL ADDRESS] : [VIRTUAL SIZE]" \
                + " : [SIZE OF DATA] : [ENTROPY]"
            tag['value'] = section['name'] + " : "
            tag['value'] += section['virtual_address'] + " : "
            tag['value'] += section['virtual_size'] + " : "
            tag['value'] += section['size_of_data'] + " : "
            tag['value'] += str(section['entropy']) 
            tags.add(pickle.dumps(tag))
    except:
        pass

    #section entropy
    try:
        for section in json_report['static']['pe_sections'] :
            tag = {}
            tag['type'] = 'section_entropy'
            tag['description'] = '[SECTION ENTROPY]'
            tag['value'] = section['entropy']
            tags.add(pickle.dumps(tag))
    except: 
        pass

    #dropped file
    try: 
        for dropped_file in json_report['dropped']:
            tag = {}
            tag['type'] = 'dropped_file'
            tag['description'] = "[NAME] : [SIZE] : [TYPE] : [MD5] : [SHA1]" \
                + " : [SHA256] : [SHA512] : [SSDEEP]"
            tag['value'] = dropped_file['name'] + " : "
            tag['value'] += str(dropped_file['size']) + " : "
            tag['value'] += dropped_file['type'] + " : "
            tag['value'] += dropped_file['md5'] + " : "
            tag['value'] += dropped_file['sha1'] + " : "
            tag['value'] += dropped_file['sha256'] + " : "
            tag['value'] += dropped_file['sha512'] + " : "
            tag['value'] += dropped_file['ssdeep'] 
            tags.add(pickle.dumps(tag))
    except:
        pass 

    #process_name
    try:
        for process in json_report['behavior']['processes'] :
            tag = {}
            tag['type'] = 'process_name'
            tag['description'] = '[PROCESS NAME]'
            tag['value'] = process['process_name']
            tags.add(pickle.dumps(tag))
    except:
        pass 

    #api 
    try:


        for process in json_report['behavior']['processes'] :
            for call in process['calls'] :
                tag = {}
                tag['type'] = 'api'
                tag['description'] = '[API CALL]'
                tag['value'] = call['api']
                tags.add(pickle.dumps(tag))  
    except:
        pass 

    #api sequence 
    try:
        #create list of all apis
        apis = list()

        for process in json_report['behavior']['processes'] :
            for call in process['calls'] :
                apis.append(call['api'])     
                  
        for i in range(0,len(apis)) :

            api_seq_string = apis[i-4] + " : " + apis[i-3] + " : " + apis[i-2]
            api_seq_string += " : " + apis[i-1] + " : " + apis[i]

            tag = {}
            tag['type'] = 'api_sequence'
            tag['description'] = '[API1] : [API2] : [API3] : [API4] : [API5]'
            tag['value'] = api_seq_string
            tags.add(pickle.dumps(tag))  
    except:
        pass 

    #file
    try:
        for file_name in json_report['behavior']['summary']['files'] :
            tag = {}
            tag['type'] = 'file'
            tag['description'] = '[FILE NAME]'
            tag['value'] = file_name
            tags.add(pickle.dumps(tag))
    except: 
        pass

    #key
    try:
        for key_name in json_report['behavior']['summary']['keys'] :
            tag = {}
            tag['type'] = 'key'
            tag['description'] = '[KEY NAME]'
            tag['value'] = key_name
            tags.add(pickle.dumps(tag))
    except: 
        pass

    #mutex
    try:
        for mutex_name in json_report['behavior']['summary']['mutexes'] :
            tag = {}
            tag['type'] = 'mutex'
            tag['description'] = '[MUTEX NAME]'
            tag['value'] = mutex_name
            tags.add(pickle.dumps(tag))
    except: 
        pass

    #events
    try:

        for event in json_report['behavior']['enhanced'] :
            tag = {}
            tag['type'] = event['event'] + "_" +event['object']
            if tag['type'] == 'load_library' or tag['type'] == 'write_file' \
                or tag['type'] == 'read_file' or tag['type'] == 'delete_file' \
                or tag['type'] == 'execute_file' or tag['type'] =='create_dir' \
                or tag['type'] =='delete_dir':       
            
                tag['description'] = '[FILE NAME]'
                tag['value'] = event['data']['file']

            elif tag['type'] == 'read_registry' or tag['type'] == \
                'delete_registry' or tag['type'] == 'write_registry' :
                
                tag['description'] = '[KEY NAME]'
                tag['value'] = event['data']['regkey']  

            elif tag['type'] == 'move_file' or tag['type'] == 'copy_file' :
                
                tag['description'] = '[FILE DEST] : [FILE SRC]'
                tag['value'] = event['data']['to'] + " : " + \
                    event['data']['from']              

            else :
                tag['description'] = '[EVENT]'
                tag['value'] = event['data']
        
            tags.add(pickle.dumps(tag))
    
    except:
        pass

    #category
    try:
        tag = {}
        tag['type'] = 'category'
        tag['description'] = '[CATEGORY]'
        tag['value'] = json_report['target']['category']
        tags.add(pickle.dumps(tag))
    except: 
        pass  

    #yara
    try:
        for yara in json_report['target']['file']['yara'] :
            tag = {}
            tag['type'] = 'yara'
            tag['description'] = '[YARA NAME] : [YARA DESCRIPTION]'
            tag['value'] = yara['name'] + " : " + yara['meta']['description']
            tags.add(pickle.dumps(tag))
    except :
        pass

    #file_type
    try:
        tag = {}
        tag['type'] = 'file_type'
        tag['description'] = '[FILE TYPE]'
        tag['value'] = json_report['target']['file']['type']
        tags.add(pickle.dumps(tag))
    except: 
        pass 
   
    #file_name
    try:
        tag = {}
        tag['type'] = 'file_name'
        tag['description'] = '[FILE NAME]'
        tag['value'] = json_report['target']['file']['name']
        tags.add(pickle.dumps(tag))
    except: 
        pass 

    #file_size
    try:
        tag = {}
        tag['type'] = 'file_size'
        tag['description'] = '[FILE SIZE]'
        tag['value'] = json_report['target']['file']['size']
        tags.add(pickle.dumps(tag))
    except: 
        pass 
    
    #AV dectection
    try:
        for AV in json_report['virustotal']['scans'] :
            if json_report['virustotal']['scans'][AV]['detected'] == True :
                tag = {}
                tag['type'] = 'av'
                tag['description'] = '[VENDOR] : [RESULT]'                
                tag['value'] = AV + " : " + \
                    json_report['virustotal']['scans'][AV]['result']
            tags.add(pickle.dumps(tag))    
    
    except :
        pass

    #udp
    try:
        for udp in json_report['network']['udp'] :
            
            udp_string = str(udp['src']) + ' : ' + str(udp['sport']) + \
                ' : ' + str(udp['dst']) + ' : ' + str(udp['dport'])

            tag = {}
            tag['type'] = 'udp'
            tag['description'] = '[SRC] : [SRC PORT] : [DST] : [DST PORT]'
            tag['value'] = udp_string
            tags.add(pickle.dumps(tag))
    except : 
        pass  
    
    #tcp
    try:

        for tcp in json_report['network']['tcp'] :
            
            tcp_string = str(tcp['src']) + ' : ' + str(tcp['sport']) + \
                ' : ' + str(tcp['dst']) + ' : ' + str(tcp['dport'])

            tag = {}
            tag['type'] = 'tcp'
            tag['description'] = '[SRC] : [SRC PORT] : [DST] : [DST PORT]'
            tag['value'] = tcp_string
            tags.add(pickle.dumps(tag))
    except : 
        pass  

    #irc
    try:
        for irc in json_report['network']['irc'] :
            tag = {}
            tag['type'] = 'irc'
            tag['description'] = '[TYPE] : [COMMAND] : [PARAMS]'
            tag['value'] = str(irc['type']) + ' : ' + str(irc['command']) + \
                ' : ' + str(irc['params'])
            tags.add(pickle.dumps(tag))
    except : 
        pass  

    #dns
    try:
        for dns in json_report['network']['dns'] :
            tag = {}
            tag['type'] = 'dns'
            tag['description'] = '[TYPE] : [REQUEST]'
            tag['value'] = str(dns['type']) + ' : ' + str(dns['request'])
            tags.add(pickle.dumps(tag))
    except : 
        pass
    
    #domain
    try:
        for domain in json_report['network']['domains'] :
            tag = {}
            tag['type'] = 'domain'
            tag['description'] = '[DOMAIN] : [IP]'
            tag['value'] = str(domain['domain']) + ' : ' + str(domain['ip'])
            tags.add(pickle.dumps(tag))
    except : 
        pass

    #icmp
    try:
        for icmp in json_report['network']['icmp'] :
            tag = {}
            tag['type'] = 'icmp'
            tag['description'] = '[SRC] : [DST] : [TYPE] : [DATA]'
            tag['value'] = str(icmp['src']) + ' : ' + str(icmp['dst']) + \
                ' : ' + str(icmp['type']) + ' : ' + str(icmp['data'])
            tags.add(pickle.dumps(tag))
    except : 
        pass

    #http_uri
    try:
        for http in json_report['network']['http'] :
            tag = {}
            tag['type'] = 'http_uri'
            tag['description'] = '[URI]'
            tag['value'] = str(http['uri'])
            tags.add(pickle.dumps(tag))
    except : 
        pass

    #http_user_agent
    try:
        for http in json_report['network']['http'] :
            tag = {}
            tag['type'] = 'http_user-agent'
            tag['description'] = '[USER AGENT]'
            tag['value'] = str(http['user-agent'])
            tags.add(pickle.dumps(tag))
    except : 
        pass

    #http_data
    try:
        for http in json_report['network']['http'] :
            tag = {}
            tag['type'] = 'http_data'
            tag['description'] = '[HTTP DATA]'
            tag['value'] = str(http['data'])
            tags.add(pickle.dumps(tag))
    except : 
        pass

    return tags

#module to save tags into the database
#params: sid - sample identifier from SID table
#tags = set of tags to add 
def store_tags(sid,tags) :

    import MySQLdb
    import _mysql_exceptions
    try :
        #username and password stored in external files
        mysql_usr = open(os.environ['HOME'] + "/.mysql_usr").read().strip()
        mysql_pwd = open(os.environ['HOME'] + "/.mysql_pwd").read().strip()
        
        #connect to database and create cursor
        db = MySQLdb.connect("localhost", mysql_usr, mysql_pwd, "mip")
        cursor = db.cursor()
    except Exception as e:
        msg = "Could not conect to db: " + e
        print msg
        logging.warning(msg)

    for tag in tags :
        try : 
            #unserialize tag    
            tag = pickle.loads(tag)
            #store into database

            s1 = MySQLdb.escape_string(tag['type'][:250]) 
            s2 = MySQLdb.escape_string(str(tag['value'])[:750])
            s3 = MySQLdb.escape_string(tag['description'])    
            #create insert command
            cmd = "INSERT INTO tags (type, value, description) VALUES (%s,%s,%s)"      
            data = (s1 , s2 , s3)
            try :     
                cursor.execute(cmd,data)
                tid = cursor.lastrowid

            #duplicate entry
            except _mysql_exceptions.IntegrityError as e :
                #parse key and value from error message
                #print e
                keyword = str(e).split("\"")[1].split()[0]          
                if keyword == 'Duplicate' :        

                    try : 
                        cmd = "SELECT tid FROM tags WHERE type=%s AND value=%s"      
                        data = (s1 , s2)                           
                        cursor.execute(cmd,data)
                        tid = cursor.fetchone()[0]
                        
                    except Exception as e:
                        logging.warning(cmd)
                        logging.warning(data)
                        logging.warning(tid)                            
                        logging.warning(e) 
                        

            #create relation in sidtid table
            try: 
                s1 = sid 
                s2 = tid
                data = (s1 , s2)        
                cmd = "INSERT INTO sidtid (SID, TID) VALUES (%s,%s)"
                cursor.execute(cmd,data)
                #print "SID/TID: " + str(sid) + "/" + str(tid) + " " + tag['type'] + " " +  str(tag['value']) + " " + tag['description']
            except Exception as e :
                pass # must pass - we don't care about duplicates at all
                        
        except Exception as e:
            print (e)
            logging.warning(e)
    
    #close and clean up connection
    try:        
        cursor.close()
        db.close()
        del cursor
        del db
    except Exception as e:
        print e
    
    return

#function to check if any errors were reported by cuckoo
def check_errors(json_report) :
    if json_report['debug']['errors'] == [] : 
        return False #no errors
    else : 
        return True #errors
